from skimage import io, transform
import os
import random
import numpy as np

class Data(object):

    '''
    label_class:标签的种类，为list
    '''

    def read_image(self, filename, batch_size, img_width = 256, img_height = 256, label_class = ['deg'], shuffle = False, if_numpy = True):
        imgs_path = []
        imgs_path = self.judge_path(imgs_path, filename)

        # 打乱训练集
        if shuffle:
            random.shuffle(imgs_path)

        batch_imgs_path = []  # 创建临时批次图片集

        for index, img_path in enumerate(imgs_path):
            batch_imgs_path.append(img_path)
            if (index+1) % batch_size == 0:
                yield (index+1) / batch_size, self.load_img(batch_imgs_path, img_width, img_height, if_numpy),\
                      self.label_all(batch_imgs_path, label_class, if_numpy)
                batch_imgs_path.clear()

    # 递归找出文件夹最里面的文件
    def judge_path(self, imgs_path, folder_path):
        if os.path.isfile(folder_path) == False:
            for sub_file in os.listdir(folder_path):
                if os.path.isfile(os.path.join(folder_path, sub_file)):
                    imgs_path.append(os.path.join(folder_path, sub_file))
                else:
                    self.judge_path(imgs_path, os.path.join(folder_path, sub_file))
        else:
            raise ValueError('The path is not a folder')
        return imgs_path

    # 图片路径加载图片
    def load_img(self, batch_imgs_path, img_width = 256, img_height = 256, if_numpy = True):
        img_set = []  # 创建图片集
        for img_path in batch_imgs_path:
            img = io.imread(img_path)
            img = transform.resize(img, (img_width, img_height))
            img_set.append(img)
        if if_numpy:
            img_set = np.asarray(img_set, np.float32)
        return img_set

    # 标记多种任务的标签（如果有的话）
    def label_all(self, batch_imgs_path, label_class, if_numpy):
        label_set = []  # 创建标签集
        for i in range(len(label_class)):
            label_set.append(list())
        for img_path in batch_imgs_path:
            for i in range(len(label_class)):
                judge_fun = getattr(self, label_class[i] + '_label')  # 拼凑字符串，得到相应的判别方法
                label_result = judge_fun(img_path)
                label_set[i].append(label_result)
        for i in range(len(label_class)):
            if if_numpy:
                label_set[i] = np.asarray(label_set[i], np.float32)
        return label_set

    # 等级标签标记
    def deg_label(self, image_path):
        if image_path.find('g0') != -1:
            label = [1, 0, 0, 0]
        elif image_path.find('g1') != -1:
            label = [0, 1, 0, 0]
        elif image_path.find('g2') != -1:
            label = [0, 0, 1, 0]
        elif image_path.find('g3') != -1:
            label = [0, 0, 0, 1]

        return label
